//! 多项式回归

use linfa::prelude::*;
use linfa_linear::LinearRegression;
use ndarray::{prelude::*};
use ndarray_rand::{rand_distr::Uniform, RandomExt};
use plotters::prelude::*;

fn main() {
    // 生成数据
    let x = Array1::random(100, Uniform::new(-3.0, 3.0));
    let y: Array1<f64> = 0.5 * &x * &x + &x + 2.0 + Array1::random(100, Uniform::new(0.0, 1.0));

    // 显示图形
    let root =
        SVGBackend::new("./examples/PolynomialRegression/data.svg", (600, 400)).into_drawing_area();
    root.fill(&WHITE).unwrap();

    let mut chat = ChartBuilder::on(&root)
        .x_label_area_size(40)
        .y_label_area_size(40)
        .build_cartesian_2d(-3.0f64..3.0f64, -2.0f64..5.0f64)
        .unwrap();

    chat.draw_series(
        x.iter()
            .zip(y.iter())
            .map(|(&x, &y)| Circle::new((x, y), 4, RED.filled())),
    )
    .unwrap();
    chat.configure_mesh()
        .disable_x_mesh()
        .disable_y_mesh()
        .draw()
        .unwrap();

    // 用线性
    let linear = linfa_linear::LinearRegression::default();
    let dataset = Dataset::new(x.to_shared().reshape((x.len(), 1)).to_owned(), y.clone());
    let model = linear.fit(&dataset).unwrap();
    println!("model: {:?}", model);
    let predict_y = model.predict(&dataset);
    println!("score: {}", predict_y.r2(&y).unwrap());
    // 画出直线
    let mut points = x
        .iter()
        .zip(predict_y.iter())
        .map(|(x, y)| (*x, *y))
        .collect::<Vec<(f64, f64)>>();
    points.sort_by(|&(x1, _), &(x2, _)| x1.partial_cmp(&x2).unwrap());
    chat.draw_series(LineSeries::new(points, BLUE)).unwrap();

    // 多项式回归的思路 增加x平方一个特征
    let mut x2 = Array2::zeros((x.len(), 0));
    x2.push_column(x.view()).unwrap();
    x2.push_column((&x * &x).view()).unwrap();
    let x2 = x2.as_standard_layout().to_owned();
    let dataset = Dataset::new(x2, y.clone());
    println!("{}", dataset.records().is_standard_layout());
    let (train, test) = dataset.split_with_ratio(0.75);
    let model = linfa_linear::LinearRegression::default();
    let model = model.fit(&train).unwrap();
    let y_predict = model.predict(&test);
    println!(" 2 score: {}", y_predict.r2(test.targets()).unwrap());

    let ret = linfa_elasticnet::ElasticNet::params().fit(&train).unwrap();
    let pre = ret.predict(&test);
    println!("e : {}", pre.r2(test.targets()).unwrap());
    let mut points = x
        .iter()
        .zip(y_predict.iter())
        .map(|(x, y)| (*x, *y))
        .collect::<Vec<_>>();
    points.sort_by(|&(x1, _), &(x2, _)| x1.partial_cmp(&x2).unwrap());
    chat.draw_series(LineSeries::new(points, GREEN)).unwrap();

    // 学习曲线
    let dataset = Dataset::new(x.to_shared().reshape((x.len(), 1)).to_owned(), y.clone());
    let (train, test) = dataset.split_with_ratio(0.75);
    println!(
        "train: {:?}, {:?}",
        train.records().shape(),
        train.targets().shape()
    );

    println!(
        "test: {:?}, {:?}",
        test.records().shape(),
        test.records().shape()
    );

    let mut train_scores = vec![];
    let mut test_scores = vec![];
    for i in 1..train.records().len_of(Axis(0)) + 1 {
        let lin_reg = LinearRegression::new();
        // 分块抽样
        let d = train.sample_chunks(i).next().unwrap();
        let model = lin_reg.fit(&d).unwrap();
        let train_y = model.predict(&d);
        let train_score = train_y.mean_squared_error(d.targets()).unwrap();
        train_scores.push(train_score);

        let y_test = model.predict(&test);
        let test_score = y_test.mean_squared_error(test.targets()).unwrap();
        test_scores.push(test_score);
    }

    // 画出来
    let _x_len = train.targets().len() + 1;

    // chat.draw_series(LineSeries::new(
    //     (1..x_len).map(|x| x as f64).zip(train_scores),
    //     BLACK,
    // ))
    // .unwrap();

    // chat.draw_series(LineSeries::new(
    //     (1..x_len).map(|x| x as f64).zip(test_scores),
    //     CYAN,
    // ))
    // .unwrap();
    println!("岭回归");
    let dataset = Dataset::new(x.to_shared().reshape((x.len(), 1)).to_owned(), y.clone());
    // let ling = linfa_linear::TweedieRegressor::params().fit(&dataset);
    let iso = linfa_linear::IsotonicRegression::default()
        .fit(&dataset)
        .unwrap();
    let pre = iso.predict(&test);
    let score = pre.mean_squared_error(test.targets()).unwrap();

    println!("score: {}", score);
    println!("score 2: {}", pre.r2(test.targets()).unwrap());
    let mut points = x
        .iter()
        .zip(pre.iter())
        .map(|(x, y)| (*x, *y))
        .collect::<Vec<_>>();
    points.sort_by(|&(x1, _), &(x2, _)| x1.partial_cmp(&x2).unwrap());
    chat.draw_series(LineSeries::new(points, BLACK)).unwrap();
}
